package com.yuba.common.ai.utils;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Flux;

/**
 * Ai聊天工具类
 */
public class AiChatUtil {

    public static Flux<ServerSentEvent<String>> chatStream(ChatClient client, String input) {
        return client
            .prompt()
            .user(input)
            .stream()
            .content()
            .map(chatResponse -> ServerSentEvent.<String>builder().data(chatResponse).build());

    }

    private static void useVectorStore(ChatClient.AdvisorSpec advisorSpec, VectorStore vectorStore) {
        // question_answer_context是一个占位符，会替换成向量数据库中查询到的文档。QuestionAnswerAdvisor会替换。
        String promptWithContext = """
            下面是上下文信息
            ---------------------
            {question_answer_context}
            ---------------------
            给定的上下文和提供的历史信息 ，而不是事先的知识，回复用户的意见。如果答案不在上下文中，就回答"暂时没有这部分知识"。
            """;
        advisorSpec.advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults(), promptWithContext));
    }

    /**
     * 构建ChatClient
     *
     * @param builder ChatClient构造器
     * @param chatModel chat模型
     * @param vectorStore 系列数据库
     * @return chatClient
     */
    public static ChatClient buildClient(ChatClient.Builder builder, ChatModel chatModel, VectorStore vectorStore) {
        return builder
            .defaultSystem("""
                你是一个渔吧知识库小助手并且是一个资深的钓鱼专家，可以为用户解决一切钓鱼相关的问题和描述钓鱼知识。
                如果用户向你打招呼，无需查询知识库，同样回复渔吧知识助手介绍
                """)
            .defaultAdvisors(advisorSpec -> {
                // 调用知识库
                useVectorStore(advisorSpec, vectorStore);
            })
            .build();
    }
}
